#include "common_lib/transform.h"

using namespace common_lib;
Transform::Transform()
{
}

Transform::~Transform()
{
}

void Transform::EulerToQuaternion(float *euler, Quaternion *quat)
{
    // init the Quaternion
    float sinThetaXOver2, cosThetaXOver2; // roll
    float sinThetaYOver2, cosThetaYOver2; // pitch
    float sinThetaZOver2, cosThetaZOver2; // yaw
    float Quaternion[4];
    // Divide the angle by two (the angles used in computing the quaternion
    //   are theta/2)

    // cout<< "euler to quaternion angle is:"<<euler[ROLL]<<","<<euler[PITCH]<<","<<euler[YAW]<<endl;

    float ThetaXOver2 = (float)0.5 * euler[ROLL];
    float ThetaYOver2 = (float)0.5 * euler[PITCH];
    float ThetaZOver2 = (float)0.5 * euler[YAW];

    sinThetaXOver2 = sin(ThetaXOver2);
    cosThetaXOver2 = cos(ThetaXOver2);
    sinThetaYOver2 = sin(ThetaYOver2);
    cosThetaYOver2 = cos(ThetaYOver2);
    sinThetaZOver2 = sin(ThetaZOver2);
    cosThetaZOver2 = cos(ThetaZOver2);

    quat->q0 = sinThetaXOver2 * sinThetaYOver2 * sinThetaZOver2 +
               cosThetaXOver2 * cosThetaYOver2 * cosThetaZOver2;

    quat->q1 = sinThetaXOver2 * cosThetaYOver2 * cosThetaZOver2 -
               cosThetaXOver2 * sinThetaYOver2 * sinThetaZOver2;

    quat->q2 = sinThetaXOver2 * cosThetaYOver2 * sinThetaZOver2 +
               cosThetaXOver2 * sinThetaYOver2 * cosThetaZOver2;

    quat->q3 = cosThetaXOver2 * cosThetaYOver2 * sinThetaZOver2 -
               sinThetaXOver2 * sinThetaYOver2 * cosThetaZOver2;

    QuatNormalize(quat);
}

void Transform::QuatNormalize(Quaternion *quat)
{
    float quatSquared[4];
    float quatMag;
    float temp;

    // Square the components of the quaternion (0 <= sSquared <= 1)
    quatSquared[Q0] = quat->q0 * quat->q0;
    quatSquared[Q1] = quat->q1 * quat->q1;
    quatSquared[Q2] = quat->q2 * quat->q2;
    quatSquared[Q3] = quat->q3 * quat->q3;

    // Find the RSS of the quaternion: sqrt( q1^2 + q2^2 + q3^2 + q4^2 )
    quatMag = quatSquared[Q0] +
              quatSquared[Q1] +
              quatSquared[Q2] +
              quatSquared[Q3];
    quatMag = sqrtf(quatMag);

    // Normalize the quaternion
    temp = 1 / quatMag;
    quat->q0 = quat->q0 * temp;
    quat->q1 = quat->q1 * temp;
    quat->q2 = quat->q2 * temp;
    quat->q3 = quat->q3 * temp;

    // #ifdef FORCE_Q0_POSITIVE
    //  Force Q0 to be positive
    if (quat->q0 < 0.0)
    {
        // #ifdef DISPLAY_DIAGNOSTIC_MSG
        //         TimingVars_DiagnosticMsg("QuatNormalize: flip sign on q");
        // #endif

        // Flip signs on all quaternion elements
        quat->q0 = -quat->q0;
        quat->q1 = -quat->q1;
        quat->q2 = -quat->q2;
        quat->q3 = -quat->q3;
    }
    // #endif
}

// Quaternion to euler
void Transform::QuaternionToEuler(Quaternion *quat, float *euler)
{
    float R[3][3];

    float q0Sq, q0q1, q0q2, q0q3;
    float q1Sq, q1q2, q1q3;
    float q2Sq, q2q3;
    float q3Sq;
    float Quaternion[4];

    Quaternion[0] = quat->q0;
    Quaternion[1] = quat->q1;
    Quaternion[2] = quat->q2;
    Quaternion[3] = quat->q3;

    // Compute values used repeatedly in the function
    q0Sq = Quaternion[Q0] * Quaternion[Q0];
    q0q1 = Quaternion[Q0] * Quaternion[Q1];
    q0q2 = Quaternion[Q0] * Quaternion[Q2];
    q0q3 = Quaternion[Q0] * Quaternion[Q3];

    q1Sq = Quaternion[Q1] * Quaternion[Q1];
    q1q2 = Quaternion[Q1] * Quaternion[Q2];
    q1q3 = Quaternion[Q1] * Quaternion[Q3];

    q2Sq = Quaternion[Q2] * Quaternion[Q2];
    q2q3 = Quaternion[Q2] * Quaternion[Q3];

    q3Sq = Quaternion[Q3] * Quaternion[Q3];

    // Form the direction cosine matrix (DCM) from q  Cbn
    R[0][0] = q0Sq + q1Sq - q2Sq - q3Sq;
    //    R[0][1] = 2.0 * (q1q2 - q0q3);
    //    R[0][2] = 2.0 * (q1q3 + q0q2);

    R[1][0] = (float)2.0 * (q1q2 + q0q3);
    //    R[1][1] = q0Sq - q1Sq + q2Sq - q3Sq;
    //    R[1][2] = 2.0 * (q2q3 - q0q1);

    R[2][0] = (float)2.0 * (q1q3 - q0q2);
    R[2][1] = (float)2.0 * (q2q3 + q0q1);
    R[2][2] = q0Sq - q1Sq - q2Sq + q3Sq;

    // Calculate the euler angles from the DCM
    // prevent NaN due to numerical error, dxg, 20180802
    if (R[2][0] > 1.0)
    {
        R[2][0] = 1.0;
    }
    if (R[2][0] < -1.0)
    {
        R[2][0] = -1.0;
    }

    euler[ROLL] = atan2(R[2][1], R[2][2]);
    euler[PITCH] = -asin(R[2][0]);
    euler[YAW] = atan2(R[1][0], R[0][0]);
}

// init Cbn from euler angles
void Transform::EulerToCbn(float *euler, Eigen::Matrix3f &cbn)
{
    float sinThetaX, cosThetaX;
    float sinThetaY, cosThetaY;
    float sinThetaZ, cosThetaZ;

    // Precompute sin/cos values used in the expressions below

    sinThetaX = sin(euler[ROLL]);
    cosThetaX = cos(euler[ROLL]);
    sinThetaY = sin(euler[PITCH]);
    cosThetaY = cos(euler[PITCH]);
    sinThetaZ = sin(euler[YAW]);
    cosThetaZ = cos(euler[YAW]);

    cbn(0, 0) = cosThetaZ * cosThetaY;
    cbn(0, 1) = cosThetaZ * sinThetaY * sinThetaX - sinThetaZ * cosThetaX;
    cbn(0, 2) = cosThetaZ * sinThetaY * cosThetaX + sinThetaZ * sinThetaX;

    cbn(1, 0) = sinThetaZ * cosThetaY;
    cbn(1, 1) = sinThetaZ * sinThetaY * sinThetaX + cosThetaZ * cosThetaX;
    cbn(1, 2) = sinThetaZ * sinThetaY * cosThetaX - cosThetaZ * sinThetaX;

    cbn(2, 0) = -sinThetaY;
    cbn(2, 1) = cosThetaY * sinThetaX;
    cbn(2, 2) = cosThetaY * cosThetaX;
}

// IMU install tilt frame to body horizon frame
void Transform::EulerToCtb(float *euler, Eigen::Matrix3f &ctb)
{
    float sinThetaX, cosThetaX;
    float sinThetaY, cosThetaY;
    float sinThetaZ, cosThetaZ;

    // Precompute sin/cos values used in the expressions below

    sinThetaX = sin(euler[ROLL]);
    cosThetaX = cos(euler[ROLL]);
    sinThetaY = sin(euler[PITCH]);
    cosThetaY = cos(euler[PITCH]);
    sinThetaZ = sin(euler[YAW]);
    cosThetaZ = cos(euler[YAW]);

    ctb(0, 0) = cosThetaZ * cosThetaY;
    ctb(0, 1) = cosThetaZ * sinThetaY * sinThetaX - sinThetaZ * cosThetaX;
    ctb(0, 2) = cosThetaZ * sinThetaY * cosThetaX + sinThetaZ * sinThetaX;

    ctb(1, 0) = sinThetaZ * cosThetaY;
    ctb(1, 1) = sinThetaZ * sinThetaY * sinThetaX + cosThetaZ * cosThetaX;
    ctb(1, 2) = sinThetaZ * sinThetaY * cosThetaX - cosThetaZ * sinThetaX;

    ctb(2, 0) = -sinThetaY;
    ctb(2, 1) = cosThetaY * sinThetaX;
    ctb(2, 2) = cosThetaY * cosThetaX;
}

// translate body to Horizon
void Transform::EulerToCbh(float *euler, Eigen::Matrix3f &cbh)
{
    float sinThetaX, cosThetaX;
    float sinThetaY, cosThetaY;
    float sinThetaZ, cosThetaZ;

    // Precompute sin/cos values used in the expressions below

    sinThetaX = sin(euler[ROLL]);
    cosThetaX = cos(euler[ROLL]);
    sinThetaY = sin(euler[PITCH]);
    cosThetaY = cos(euler[PITCH]);
    sinThetaZ = sin(0);
    cosThetaZ = cos(0);

    cbh(0, 0) = cosThetaZ * cosThetaY;
    cbh(0, 1) = cosThetaZ * sinThetaY * sinThetaX - sinThetaZ * cosThetaX;
    cbh(0, 2) = cosThetaZ * sinThetaY * cosThetaX + sinThetaZ * sinThetaX;

    cbh(1, 0) = sinThetaZ * cosThetaY;
    cbh(1, 1) = sinThetaZ * sinThetaY * sinThetaX + cosThetaZ * cosThetaX;
    cbh(1, 2) = sinThetaZ * sinThetaY * cosThetaX - cosThetaZ * sinThetaX;

    cbh(2, 0) = -sinThetaY;
    cbh(2, 1) = cosThetaY * sinThetaX;
    cbh(2, 2) = cosThetaY * cosThetaX;
}

Eigen::Matrix<float, 6, 1> Transform::GetXyzyprFromTransformMatrix(const Eigen::Matrix4f &T)
{
    Eigen::Matrix<float, 6, 1> xyz_ypr = Eigen::Matrix<float, 6, 1>::Zero();

    xyz_ypr.block(0, 0, 3, 1) = T.block(0, 3, 3, 1);

    Eigen::Matrix3f rotation_matrix = T.block<3, 3>(0, 0);
    Eigen::Vector3f euler_rpy = GetEulerFromRotationMatrix(rotation_matrix);

    xyz_ypr(3) = euler_rpy.z();
    xyz_ypr(4) = euler_rpy.y();
    xyz_ypr(5) = euler_rpy.x();

    return xyz_ypr;
}

Eigen::Vector3f Transform::GetEulerFromRotationMatrix(Eigen::Matrix3f rotation_matrix)
{
    Eigen::Vector3f euler_rpy;
    euler_rpy.x() = atan2(rotation_matrix(2, 1), rotation_matrix(2, 2));
    euler_rpy.y() = -asin(rotation_matrix(2, 0));
    euler_rpy.z() = atan2(rotation_matrix(1, 0), rotation_matrix(0, 0));
    return euler_rpy;
}

Eigen::Matrix3d Transform::EulerAngle2Mat(double rx, double ry, double rz)
{
    Eigen::AngleAxisd rollAngle(Eigen::AngleAxisd(rx, Eigen::Vector3d::UnitX()));
    Eigen::AngleAxisd pitchAngle(Eigen::AngleAxisd(ry, Eigen::Vector3d::UnitY()));
    Eigen::AngleAxisd yawAngle(Eigen::AngleAxisd(rz, Eigen::Vector3d::UnitZ()));
    Eigen::Matrix3d rotation_matrix;
    rotation_matrix = yawAngle * pitchAngle * rollAngle;
    return rotation_matrix;
}

Eigen::Matrix3d Transform::EulerAngle2Mat(Eigen::Vector3d eular)
{
    return EulerAngle2Mat(eular.x(), eular.y(), eular.z());
}

Eigen::Matrix4d Transform::Compose(Eigen::Vector3d positon, Eigen::Vector3d rotEular)
{
    Eigen::Matrix3d rot = EulerAngle2Mat(rotEular);
    Eigen::Matrix4d t;
    t.setIdentity();
    t.block<3, 3>(0, 0) = rot;
    t.block<3, 1>(0, 3) = positon;
    return t;
}

Eigen::Matrix4d Transform::ComposeEuler(const double x, const double y,
                                        const double z, const double rx,
                                        const double ry, const double rz)
{
    Eigen::Vector3d rot(rx, ry, rz);
    Eigen::Vector3d pos(x, y, z);
    return Compose(pos, rot);
}

Eigen::Vector3d Transform::Mat2EulerAngle(Eigen::Matrix3d mat)
{
    Eigen::Vector3d rot = mat.eulerAngles(2, 1, 0);
    rot = rot / M_PI * 180;
    return rot;
}

Eigen::VectorXd Transform::H2EulerAngle(Eigen::Matrix4d t)
{
    Eigen::VectorXd pose = Eigen::VectorXd(6);
    Eigen::Matrix3d mt = t.block<3, 3>(0, 0);
    Eigen::Vector3d p3 = t.block<3, 1>(0, 3).col(0);
    pose(0, 0) = p3.x();
    pose(1, 0) = p3.y();
    pose(2, 0) = p3.z();
    Eigen::Vector3d eular = Mat2EulerAngle(mt);
    pose(3, 0) = eular.x();
    pose(4, 0) = eular.y();
    pose(5, 0) = eular.z();
    return pose;
}

Eigen::Vector3d Transform::PosTF(Eigen::Vector3d &input, geometry_msgs::Pose2D &pose)
{
    Eigen::Vector3d ret;
    Eigen::Vector3d t = Eigen::Vector3d(input(0) - pose.x, input(1) - pose.y, 0.0f);
    Eigen::Quaterniond q = Eigen::AngleAxisd(-pose.theta, Eigen::Vector3d::UnitZ()) *
                           Eigen::AngleAxisd(0.0, Eigen::Vector3d::UnitY()) *
                           Eigen::AngleAxisd(0.0, Eigen::Vector3d::UnitX());

    ret = q * t;
    return ret;
}

Eigen::Vector3d Transform::PosTF(Eigen::Vector3d &input, geometry_msgs::Pose &pose)
{
    Eigen::Vector3d ret;
    Eigen::Vector3d t = Eigen::Vector3d(input(0) - pose.position.x, input(1) - pose.position.y, 0.0f);
    Eigen::Quaterniond q = Eigen::AngleAxisd(-tf::getYaw(pose.orientation), Eigen::Vector3d::UnitZ()) *
                           Eigen::AngleAxisd(0.0, Eigen::Vector3d::UnitY()) *
                           Eigen::AngleAxisd(0.0, Eigen::Vector3d::UnitX());
    ret = q * t;
    return ret;
}